import torchaudio as ta
from chatterbox.tts import ChatterboxTTS
import time
import os

def generate_tts(text, output_path=None, audio_prompt_path=None, device="cpu", exaggeration=0.5, cfg_weight=0.5, temperature=0.8):
    """
    使用ChatterboxTTS生成语音，支持指定音频提示
    
    参数:
        text: 要合成的文本
        output_path: 输出音频文件路径，如果不指定则自动生成
        audio_prompt_path: 音频提示文件路径，如果存在则使用该声音风格
        device: 使用的设备，可选"cpu"或"cuda"
        exaggeration: 情感夸张度，0.5为正常，值越大情感越强烈
        cfg_weight: 生成控制强度，值越低语速越慢更有表现力
        temperature: 温度参数，控制随机性，越高越随机
    
    返回:
        output_path: 保存的音频文件路径
        elapsed_time: 耗时(秒)
    """
    start_time = time.time()
    
    model = ChatterboxTTS.from_pretrained(device=device)
    
    if audio_prompt_path:
        wav = model.generate(
            text, 
            audio_prompt_path=audio_prompt_path,
            exaggeration=exaggeration,
            cfg_weight=cfg_weight,
            temperature=temperature
        )
    else:
        wav = model.generate(
            text,
            exaggeration=exaggeration,
            cfg_weight=cfg_weight,
            temperature=temperature
        )
    
    # 如果未指定输出路径，则生成一个默认路径
    if output_path is None:
        prefix = "prompt_" if audio_prompt_path else "default_"
        timestamp = int(time.time())
        output_path = f"{prefix}{timestamp}.wav"
    
    # 保存音频文件
    ta.save(output_path, wav, model.sr)
    
    elapsed_time = time.time() - start_time
    return output_path, elapsed_time

if __name__ == "__main__":
    # 中文示例
    text_cn = "China No.1"
    
    # 不使用音频提示
    output_path, elapsed_time = generate_tts(
        text_cn, 
        output_path="test-1.wav",
        exaggeration=0.7,  # 增加情感表现力
        cfg_weight=0.4,  # 调整语速和表现力
        temperature=0.7  # 稍微降低随机性
    )
    print(f"生成完成，文件保存为: {output_path}，耗时: {elapsed_time:.2f}秒")
    
    # # 使用音频提示
    # AUDIO_PROMPT_PATH = "sichuan-000001.wav"
    # output_path, elapsed_time = generate_tts(
    #     text_cn, 
    #     output_path="test-2.wav", 
    #     audio_prompt_path=AUDIO_PROMPT_PATH,
    #     exaggeration=0.6,
    #     cfg_weight=0.4
    # )
    # print(f"生成完成，文件保存为: {output_path}，耗时: {elapsed_time:.2f}秒")